{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mxnet BYOM: Train locally and deploy on SageMaker.\n",
    "\n",
    "1. [Introduction](#Introduction)\n",
    "2. [Prerequisites and Preprocessing](#Prequisites-and-Preprocessing)\n",
    "    1. [Permissions and environment variables](#Permissions-and-environment-variables)\n",
    "    2. [Data Setup](#Data-setup)\n",
    "3. [Training the network locally](#Training)\n",
    "4. [Set up hosting for the model](#Set-up-hosting-for-the-model)\n",
    "    1. [Export from MXNet](#Export-the-model-from-mxnet)\n",
    "    2. [Import model into SageMaker](#Import-model-into-SageMaker)\n",
    "    3. [Create endpoint](#Create-endpoint) \n",
    "5. [Validate the endpoint for use](#Validate-the-endpoint-for-use)\n",
    "\n",
    "\n",
    "__Note__: Compare this with the [tensorflow bring your own model example](../tensorflow_iris_byom/tensorflow_BYOM_iris.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "In this notebook, we will train a neural network locally on the location from where this notebook is run using MXNet. We will then see how to create an endpoint from the trained MXNet model and deploy it on SageMaker. We will then inference from the newly created SageMaker endpoint. \n",
    "\n",
    "The neural network that we will use is a simple fully-connected neural network. The definition of the neural network can be found in the accompanying [mnist.py](mnist.py) file. The ``build_graph`` method contains the model defnition (shown below).\n",
    "\n",
    "```python\n",
    "def build_graph():\n",
    "    data = mx.sym.var('data')\n",
    "    data = mx.sym.flatten(data=data)\n",
    "    fc1 = mx.sym.FullyConnected(data=data, num_hidden=128)\n",
    "    act1 = mx.sym.Activation(data=fc1, act_type=\"relu\")\n",
    "    fc2 = mx.sym.FullyConnected(data=act1, num_hidden=64)\n",
    "    act2 = mx.sym.Activation(data=fc2, act_type=\"relu\")\n",
    "    fc3 = mx.sym.FullyConnected(data=act2, num_hidden=10)\n",
    "    return mx.sym.SoftmaxOutput(data=fc3, name='softmax')\n",
    "```\n",
    "\n",
    "From this definitnion we can see that there are two fully-connected layers of 128 and 64 neurons each. The activations of the last fully-connected layer is then fed into a Softmax layer of 10 neurons. We use 10 neurons here because the datatset on which we are going to predict is the MNIST dataset of hand-written digit recognition which has 10 classes. More details can be found about the dataset on the [creator's webpage](http://yann.lecun.com/exdb/mnist/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prequisites and Preprocessing\n",
    "\n",
    "### Permissions and environment variables\n",
    "\n",
    "Here we set up the linkage and authentication to AWS services. In this notebook we only need the roles used to give learning and hosting access to your data. The Sagemaker SDK will use S3 defualt buckets when needed. Supply the role in the variable below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "isConfigCell": true
   },
   "outputs": [],
   "source": [
    "import boto3, re\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data setup\n",
    "\n",
    "Next, we need to pull the data from the author's site to our local box. Since we have ``mxnet`` utilities, we will use the utilities to download the dataset locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "data = mx.test_utils.get_mnist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### Training\n",
    "\n",
    "It is time to train the network. Since we are training the network locally, we can make use of mxnet training tools. The training method is also in the accompanying [mnist.py](mnist.py) file. The notebook assumes that this instance is a `p2.xlarge`. If running this in a non-GPU notebook instance, please adjust num_gpus=0 and num_cpu=1 The method is shown below. \n",
    "\n",
    "```python \n",
    "def train(data, hyperparameters= {'learning_rate': 0.11}, num_cpus=0, num_gpus =1 , **kwargs):\n",
    "    train_labels = data['train_label']\n",
    "    train_images = data['train_data']\n",
    "    test_labels = data['test_label']\n",
    "    test_images = data['test_data']\n",
    "    batch_size = 100\n",
    "    train_iter = mx.io.NDArrayIter(train_images, train_labels, batch_size, shuffle=True)\n",
    "    val_iter = mx.io.NDArrayIter(test_images, test_labels, batch_size)\n",
    "    logging.getLogger().setLevel(logging.DEBUG)\n",
    "    mlp_model = mx.mod.Module(\n",
    "        symbol=build_graph(),\n",
    "        context=get_train_context(num_cpus, num_gpus))\n",
    "    mlp_model.fit(train_iter,\n",
    "                  eval_data=val_iter,\n",
    "                  optimizer='sgd',\n",
    "                  optimizer_params={'learning_rate': float(hyperparameters.get(\"learning_rate\", 0.1))},\n",
    "                  eval_metric='acc',\n",
    "                  batch_end_callback=mx.callback.Speedometer(batch_size, 100),\n",
    "                  num_epoch=10)\n",
    "    return mlp_model\n",
    "```\n",
    "\n",
    "The method above collects the ``data`` variable that ``get_mnist`` method gives you (which is a dictionary of data arrays) along with a dictionary of ``hyperparameters`` which only contains learning rate, and other parameters. It creates a [``mxnet.mod.Module``](https://mxnet.incubator.apache.org/api/python/module.html) from the network graph we built in the ``build_graph`` method and trains the network using the ``mxnet.mod.Module.fit`` method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mnist import train\n",
    "model = train(data = data, num_cpus=0, num_gpus=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to run the training on a cpu or if you are on an instance with cpus only, pass appropriate arguments. \n",
    "\n",
    "## Set up hosting for the model\n",
    "\n",
    "### Export the model from mxnet\n",
    "\n",
    "In order to set up hosting, we have to import the model from training to hosting. We will begin by exporting the model from MXNet and saving it down. Analogous to the [TensorFlow example](../tensorflow_iris_byom/tensorflow_BYOM_iris.ipynb), some structure needs to be followed. The exported model has to be converted into a form that is readable by ``sagemaker.mxnet.model.MXNetModel``. The following code describes exporting the model in a form that does the same:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "os.mkdir('model')\n",
    "\n",
    "model.save_checkpoint('model/model', 0000)\n",
    "with open ( 'model/model-shapes.json', \"w\") as shapes:\n",
    "    json.dump([{\"shape\": model.data_shapes[0][1], \"name\": \"data\"}], shapes)\n",
    "\n",
    "import tarfile\n",
    "def flatten(tarinfo):\n",
    "    tarinfo.name = os.path.basename(tarinfo.name)\n",
    "    return tarinfo\n",
    "\n",
    "tar = tarfile.open(\"model.tar.gz\", \"w:gz\")\n",
    "tar.add(\"model\", filter=flatten)\n",
    "tar.close()    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above piece of code essentially hacks the MXNet model export into a sagemaker-readable model export. Study the exported model files if you want to organize your exports in the same fashion as well. Alternatively, you can load the model on MXNet itself and load the sagemaker model as you normally would. Refer [here](https://github.com/aws/sagemaker-python-sdk#model-loading) for details on how to load MXNet models.\n",
    "\n",
    "### Import model into SageMaker\n",
    "\n",
    "Open a new sagemaker session and upload the model on to the default S3 bucket. We can use the ``sagemaker.Session.upload_data`` method to do this. We need the location of where we exported the model from MXNet and where in our default bucket we want to store the model(``/model``). The default S3 bucket can be found using the ``sagemaker.Session.default_bucket`` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "inputs = sagemaker_session.upload_data(path='model.tar.gz', key_prefix='model')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the ``sagemaker.mxnet.model.MXNetModel`` to import the model into SageMaker that can be deployed. We need the location of the S3 bucket where we have the model, the role for authentication and the entry_point where the model defintion is stored (``mnist.py``). The import call is the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.mxnet.model import MXNetModel\n",
    "sagemaker_model = MXNetModel(model_data = 's3://' + sagemaker_session.default_bucket() + '/model/model.tar.gz',\n",
    "                                  role = role,\n",
    "                                  entry_point = 'mnist.py')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create endpoint\n",
    "\n",
    "Now the model is ready to be deployed at a SageMaker endpoint. We can use the ``sagemaker.mxnet.model.MXNetModel.deploy`` method to do this. Unless you have created or prefer other instances, we recommend using 1 ``'ml.c4.xlarge'`` instance for this training. These are supplied as arguments. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "logging.getLogger().setLevel(logging.WARNING)\n",
    "\n",
    "predictor = sagemaker_model.deploy(initial_instance_count=1,\n",
    "                                          instance_type='ml.m4.xlarge')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Validate the endpoint for use\n",
    "\n",
    "We can now use this endpoint to classify hand-written digits. To see inference in action, draw a digit in the image box below. The pixel data from your drawing will be loaded into a ``data`` variable in this notebook. \n",
    "\n",
    "*Note: after drawing the image, you'll need to move to the next notebook cell.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "HTML(open(\"input.html\").read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = predictor.predict(data)\n",
    "print('Raw prediction result:')\n",
    "print(response)\n",
    "\n",
    "labeled_predictions = list(zip(range(10), response[0]))\n",
    "print('Labeled predictions: ')\n",
    "print(labeled_predictions)\n",
    "\n",
    "labeled_predictions.sort(key=lambda label_and_prob: 1.0 - label_and_prob[1])\n",
    "print('Most likely answer: {}'.format(labeled_predictions[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "(Optional) Delete the Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(predictor.endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you do not want continied use of the endpoint, you can remove it. Remember, open endpoints are charged. If this is a simple test or practice, it is recommended to delete them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sagemaker.Session().delete_endpoint(predictor.endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clear all stored model data so that we don't overwrite them the next time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.remove('model.tar.gz')\n",
    "import shutil\n",
    "shutil.rmtree('model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_p36",
   "language": "python",
   "name": "conda_mxnet_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  },
  "notice": "Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.  Licensed under the Apache License, Version 2.0 (the \"License\"). You may not use this file except in compliance with the License. A copy of the License is located at http://aws.amazon.com/apache2.0/ or in the \"license\" file accompanying this file. This file is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
